from numpy import random, matrix, power
from scipy import *
from pylab import grid, show, plot

def get_cov_matrix(ro, sigma0, sigma1):
   '''
       ro = cov(x, y) / (sqrt(Dx) * sqrt(Dy))
       cov(x, y) = ro * sigma0 * sigma1
   '''
   m = zeros((2,2), dtype=double)
   if((ro <=1 and ro >= (-1)) and (sigma0 > 0) and (sigma1 > 0)):
       m[0,0] = power(sigma0, 2)
       m[0,1] = ro * (sigma0 * sigma1)
       m[1,0] = m[0,1]
       m[1,1] = power(sigma1, 2)
   else:
       pass
   return m


def jointnorm_rvs(mu0, mu1, sigma0, sigma1, ro):
   '''
       generate pairs of random variables follow
       joint normal distribution
   '''
   m_cov = get_cov_matrix(ro, sigma0, sigma1)
   return random.multivariate_normal([mu0, mu1], m_cov)

def test_normal(mu=0, sigma=1):
   return stats.norm.rvs(mu, sigma)

def test_main():

   print get_cov_matrix(0.8, 0.4, 0.4)
   #print get_cov_matrix(0.0, 0.4, 0.4)
   #iteration=2500
   #v1 = zeros(iteration, dtype=double)
   #v2 = zeros(iteration, dtype=double)
   #for i in xrange(0, iteration):
       #temp = jointnorm_rvs(0.5, 0.5, 0.4, 0.4, 0.8)
       #v1[i] = temp[0]
       #v2[i] = temp[1]
       #v1[i] = test_normal()
   #grid(True)
   #plot(v1, '.r')
   #show()

   pass

if __name__ == '__main__':
   test_main()
